import numpy as np

"""
马尔科夫转移矩阵试验
案例来源：https://zhuanlan.zhihu.com/p/26453269

马尔科夫链公式概括：
转移矩阵 A： n x n维，其中行向量各元素和=1
状态矩阵 P： 1 x n维
初始状态为P<0>
P<i>表示第n期的状态
则
P<i> = P<i-1> 点乘 A 或 P<i> = P<0> 点乘 A^n
状态矩阵P<i>只依赖于其前一刻的状态P<i-1>，和其他期数的状态无关

一个啥子王二狗，每天12点时，只有三种状态：吃、玩、睡
每种状态的转移矩阵是
    [
        [0.2, 0.3, 0.5], 
        [0.1, 0.6, 0.3],
        [0.4, 0.5, 0.1]
    ]

如果当天初始状态概率为[0.6, 0.2, 0.2]，
求T+2的状态是什么？
"""


def get_status_at_n(init_status, trans_mat, n):
    shift = trans_mat
    if n > 1:
        for i in range(1, n):
            shift = shift.dot(trans_mat)
            print("shift:", shift)
    status = init_status.dot(shift)
    return status


# 转移矩阵
trans_mat = np.array([
    # 吃 玩 睡
    [0.2, 0.3, 0.5],
    [0.1, 0.6, 0.3],
    [0.4, 0.5, 0.1],
])

# 初始状态
init_status = np.array([[0.6, 0.2, 0.2]])

# 求解
n = 2
r = get_status_at_n(init_status, trans_mat, n)
print(r)

